{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1223e5cc",
   "metadata": {},
   "source": [
    "# 스트리밍 모드"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81983922",
   "metadata": {},
   "source": [
    "`graph`의 전체 상태를 **스트리밍**하는 방법"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b282c6",
   "metadata": {},
   "source": [
    "LangGraph는 여러 스트리밍(Streaming) 모드를 지원합니다. \n",
    "\n",
    "주요 모드는 다음과 같습니다\n",
    "\n",
    "- `values`: 이 스트리밍 모드는 그래프의 값들을 스트리밍합니다. 이는 각 노드가 호출된 후의 **그래프의 전체 상태**를 의미합니다.\n",
    "- `updates`: 이 스트리밍 모드는 그래프의 업데이트 내용을 스트리밍합니다. 이는 각 노드가 호출된 후의 **그래프 상태에 대한 업데이트**를 의미합니다.\n",
    "- `messages`: 이 스트리밍 모드는 각 노드의 메시지를 스트리밍합니다. 이때 **LLM 의 토큰 단위의 출력 스트리밍** 도 가능합니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "976c957d",
   "metadata": {},
   "source": [
    "## 환경설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a48be045",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf34145",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74df71b5",
   "metadata": {},
   "source": [
    "## 그래프 정의하기\n",
    "\n",
    "이 가이드에서는 간단한 에이전트를 사용하겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "540e6b63",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List, Dict\n",
    "from typing_extensions import TypedDict\n",
    "from langchain.tools import tool\n",
    "from langchain_teddynote.tools import GoogleNews\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "\n",
    "\n",
    "########## 1. 상태 정의 ##########\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    # 메시지 목록 주석 추가\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "########## 2. 도구 정의 및 바인딩 ##########\n",
    "# 도구 초기화\n",
    "# 키워드로 뉴스 검색하는 도구 생성\n",
    "@tool\n",
    "def search_keyword(query: str) -> List[Dict[str, str]]:\n",
    "    \"\"\"Look up news by keyword\"\"\"\n",
    "    news_tool = GoogleNews()\n",
    "    return \"\\n\".join(\n",
    "        [f'- {news[\"content\"]}' for news in news_tool.search_by_keyword(query, k=5)]\n",
    "    )\n",
    "\n",
    "\n",
    "# 도구 리스트 생성\n",
    "tools = [search_keyword]\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# 도구와 LLM 결합\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "\n",
    "########## 3. 노드 추가 ##########\n",
    "# 챗봇 함수 정의\n",
    "def chatbot(state: State):\n",
    "    # 메시지 호출 및 반환\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "\n",
    "# 상태 그래프 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 챗봇 노드 추가\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "# 도구 노드 생성 및 추가\n",
    "tool_node = ToolNode(tools=[search_keyword])\n",
    "\n",
    "# 도구 노드 추가\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "# 조건부 엣지\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"chatbot\",\n",
    "    tools_condition,\n",
    ")\n",
    "\n",
    "########## 4. 엣지 추가 ##########\n",
    "\n",
    "# tools > chatbot\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "# START > chatbot\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# chatbot > END\n",
    "graph_builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "# 그래프 컴파일\n",
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2beb8ce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "visualize_graph(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c8c3375",
   "metadata": {},
   "source": [
    "## 노드의 단계별 출력\n",
    "\n",
    "**스트리밍 모드**\n",
    "- `values`: 각 단계의 현재 상태 값 출력\n",
    "- `updates`: 각 단계의 상태 업데이트만 출력 (기본값)\n",
    "- `messages`: 각 단계의 메시지 출력\n",
    "\n",
    "여기서 스트리밍의 의미는 LLM 출력시 토큰 단위로 스트리밍하는 개념이 아니라, 단계별로 출력하는 의미를 가집니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "418bf13d",
   "metadata": {},
   "source": [
    "### `stream_mode = \"values\"`\n",
    "\n",
    "`values` 모드는 각 단계의 현재 상태 값을 출력합니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "`chunk.items()`\n",
    "\n",
    "- `key`: State 의 key 값\n",
    "- `value`: State 의 key 에 대한하는 value"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da92a102",
   "metadata": {},
   "source": [
    "#### 동기(Synchronous) 방식의 스트리밍\n",
    "\n",
    "- `chunk` 는 dictionary 형태(key: State 의 key, value: State 의 value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09a98b5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 동기 스트림 처리(stream_mode=\"values\")\n",
    "for chunk in graph.stream(inputs, stream_mode=\"values\"):\n",
    "\n",
    "    # chunk 는 dictionary 형태(key: State 의 key, value: State 의 value)\n",
    "    for state_key, state_value in chunk.items():\n",
    "        if state_key == \"messages\":\n",
    "            state_value[-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3451768c",
   "metadata": {},
   "source": [
    "#### 비동기(Asynchronous) 방식의 스트리밍\n",
    "\n",
    "**참고**\n",
    "\n",
    "- `astream()` 메서드는 비동기 스트림 처리를 통해 그래프를 실행하고 값 모드로 청크 단위 응답을 생성합니다.\n",
    "- `async for` 문을 사용하여 비동기 스트림 처리를 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b5ec20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 비동기 스트림 처리(stream_mode=\"values\")\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"values\"):\n",
    "    # chunk 는 dictionary 형태(key: State 의 key, value: State 의 value)\n",
    "    for state_key, state_value in chunk.items():\n",
    "        if state_key == \"messages\":\n",
    "            state_value[-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95a348b2",
   "metadata": {},
   "source": [
    "최종 결과만 확인하고 싶은 경우, 다음과 같이 처리 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde0bb39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "final_result = None\n",
    "\n",
    "# 비동기 스트림 처리(stream_mode=\"values\")\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"values\"):\n",
    "    final_result = chunk\n",
    "\n",
    "# 최종 결과 출력\n",
    "print(final_result[\"messages\"][-1].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bf6277c",
   "metadata": {},
   "source": [
    "### `stream_mode = \"updates\"`\n",
    "\n",
    "`updates` 모드는 각 단계에 대한 업데이트된 State 만 내보냅니다. \n",
    "\n",
    "- 출력은 노드 이름을 key 로, 업데이트된 값을 values 으로 하는 `dictionary` 입니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "`chunk.items()`\n",
    "\n",
    "- `key`: 노드(Node) 의 이름\n",
    "- `value`: 해당 노드(Node) 단계에서의 출력 값(dictionary). 즉, 여러 개의 key-value 쌍을 가진 dictionary 입니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b65b96d7",
   "metadata": {},
   "source": [
    "#### 동기(Synchronous) 방식의 스트리밍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78ef038e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 동기 스트림 처리(stream_mode=\"updates\")\n",
    "for chunk in graph.stream(inputs, stream_mode=\"updates\"):\n",
    "    # chunk 는 dictionary 형태(key: 노드, value: 노드의 상태 값)\n",
    "    for node, value in chunk.items():\n",
    "        if node:\n",
    "            print(f\"\\n[Node: {node}]\\n\")\n",
    "        if \"messages\" in value:\n",
    "            value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad4b30b",
   "metadata": {},
   "source": [
    "#### 비동기(Asynchronous) 방식의 스트리밍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b299547",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 비동기 스트림 처리(stream_mode=\"updates\")\n",
    "async for chunk in graph.astream(inputs, stream_mode=\"updates\"):\n",
    "    # chunk 는 dictionary 형태(key: 노드, value: 노드의 상태 값)\n",
    "    for node, value in chunk.items():\n",
    "        if node:\n",
    "            print(f\"\\n[Node: {node}]\\n\")\n",
    "        if \"messages\" in value:\n",
    "            value[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41832f7a",
   "metadata": {},
   "source": [
    "### `stream_mode = \"messages\"`\n",
    "\n",
    "`messages` 모드는 각 단계에 대한 메시지를 스트리밍합니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- `chunk` 는 두 개의 요소를 가진 tuple 입니다.\n",
    "  - `chunk_msg`: 실시간 출력 메시지\n",
    "  - `metadata`: 노드 정보"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df033eff",
   "metadata": {},
   "source": [
    "#### 동기(Synchronous) 방식의 스트리밍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e846ef38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 동기 스트림 처리(stream_mode=\"messages\")\n",
    "# chunk_msg: 실시간 출력 메시지, metadata: 노드 정보\n",
    "for chunk_msg, metadata in graph.stream(inputs, stream_mode=\"messages\"):\n",
    "\n",
    "    # chatbot 노드에서 출력된 메시지만 출력\n",
    "    if metadata[\"langgraph_node\"] == \"chatbot\":\n",
    "        if chunk_msg.content:\n",
    "            print(chunk_msg.content, end=\"\", flush=True)\n",
    "\n",
    "    else:\n",
    "        print(chunk_msg.content)\n",
    "        print(f\"\\n\\nmetadata: \\n{metadata}\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7c4fdef",
   "metadata": {},
   "source": [
    "#### 비동기(Asynchronous) 방식의 스트리밍"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4186d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 비동기 스트림 처리(stream_mode=\"messages\")\n",
    "# chunk_msg: 실시간 출력 메시지, metadata: 노드 정보\n",
    "async for chunk_msg, metadata in graph.astream(inputs, stream_mode=\"messages\"):\n",
    "    # chatbot 노드에서 출력된 메시지만 출력\n",
    "    if metadata[\"langgraph_node\"] == \"chatbot\":\n",
    "        if chunk_msg.content:\n",
    "            print(chunk_msg.content, end=\"\", flush=True)\n",
    "    else:\n",
    "        print(chunk_msg.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e08f9ac",
   "metadata": {},
   "source": [
    "## 특정 노드에 대한 출력 스트리밍\n",
    "\n",
    "**참고**\n",
    "\n",
    "- `metadata[\"langgraph_node\"]` 를 통해 특정 노드에서 출력된 메시지만 출력할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ea8d07e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List, Dict\n",
    "from typing_extensions import TypedDict\n",
    "from langchain.tools import tool\n",
    "from langchain_teddynote.tools import GoogleNews\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "\n",
    "\n",
    "########## 1. 상태 정의 ##########\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    # 메시지 목록 주석 추가\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "########## 2. 도구 정의 및 바인딩 ##########\n",
    "# 도구 초기화\n",
    "# 키워드로 뉴스 검색하는 도구 생성\n",
    "@tool\n",
    "def search_keyword(query: str) -> List[Dict[str, str]]:\n",
    "    \"\"\"Look up news by keyword\"\"\"\n",
    "    news_tool = GoogleNews()\n",
    "    return \"\\n\".join(\n",
    "        [f'- {news[\"content\"]}' for news in news_tool.search_by_keyword(query, k=5)]\n",
    "    )\n",
    "\n",
    "\n",
    "# 도구 리스트 생성\n",
    "tools = [search_keyword]\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# 도구와 LLM 결합 (tags 추가)\n",
    "llm_with_tools = llm.bind_tools(tools).with_config(tags=[\"WANT_TO_STREAM\"])\n",
    "\n",
    "\n",
    "########## 3. 노드 추가 ##########\n",
    "# 챗봇 함수 정의\n",
    "def chatbot(state: State):\n",
    "    # 메시지 호출 및 반환\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "\n",
    "# 상태 그래프 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 챗봇 노드 추가\n",
    "graph_builder.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "# 도구 노드 생성 및 추가\n",
    "tool_node = ToolNode(tools=[search_keyword])\n",
    "\n",
    "# 도구 노드 추가\n",
    "graph_builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "# 조건부 엣지\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"chatbot\",\n",
    "    tools_condition,\n",
    ")\n",
    "\n",
    "########## 4. 엣지 추가 ##########\n",
    "\n",
    "# tools > chatbot\n",
    "graph_builder.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "# START > chatbot\n",
    "graph_builder.add_edge(START, \"chatbot\")\n",
    "\n",
    "# chatbot > END\n",
    "graph_builder.add_edge(\"chatbot\", END)\n",
    "\n",
    "# 그래프 컴파일\n",
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "610cdf45",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_graph(graph, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb47cbe3",
   "metadata": {},
   "source": [
    "특정 노드(Node) 에 대해서 출력하고 싶은 경우, stream_mode=\"messages\" 를 통해 설정할 수 있습니다.\n",
    "\n",
    "`stream_mode=\"messages\"` 설정시, (`chunk_msg`, `metadata`) 형태로 메시지를 받습니다.\n",
    "- `chunk_msg` 는 실시간 출력 메시지, \n",
    "- `metadata` 는 노드 정보를 의미합니다.\n",
    "\n",
    "`metadata[\"langgraph_node\"]` 를 통해 특정 노드에서 출력된 메시지만 출력할 수 있습니다.\n",
    "\n",
    "(예시) chatbot 노드에서 출력된 메시지만 출력하는 경우\n",
    "\n",
    "`metadata[\"langgraph_node\"] == \"chatbot\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1eedddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "# 사용자의 메시지를 딕셔너리 형태로 입력 데이터 구성\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# stream_mode=\"messages\" 를 통한 스트리밍 처리\n",
    "for chunk_msg, metadata in graph.stream(inputs, stream_mode=\"messages\"):\n",
    "    # HumanMessage 가 아닌 최종 노드의 유효한 컨텐츠만 출력 처리\n",
    "    if (\n",
    "        chunk_msg.content\n",
    "        and not isinstance(chunk_msg, HumanMessage)\n",
    "        and metadata[\"langgraph_node\"] == \"chatbot\"\n",
    "    ):\n",
    "        print(chunk_msg.content, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84cc7d5b",
   "metadata": {},
   "source": [
    "metadata 를 출력하면 노드 정보를 확인할 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60550f25",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fc31981",
   "metadata": {},
   "source": [
    "## 사용자 정의 `tag` 필터링 된 스트리밍\n",
    "\n",
    "LLM 의 출력이 여러 군데에서 발생하는 경우, 특정 노드에서 출력된 메시지만 출력하고 싶은 경우가 있습니다.\n",
    "\n",
    "이러한 경우, `tags` 를 추가하여 출력하고 싶은 노드만 선별할 수 있습니다.\n",
    "\n",
    "llm 에 tags 를 추가하는 방법은 다음과 같습니다. `tags` 는 리스트 형태로 추가할 수 있습니다.\n",
    "\n",
    "`llm.with_config(tags=[\"WANT_TO_STREAM\"])`\n",
    "\n",
    "\n",
    "이를 통해 이벤트를 더 정확하게 필터링하여 해당 모델에서 발생한 이벤트만 유지할 수 있습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbfc5067",
   "metadata": {},
   "source": [
    "아래 예시는 `WANT_TO_STREAM` 태그가 있는 경우만 출력하는 예시입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b53e958",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 비동기 이벤트 스트림 처리(astream_events)\n",
    "async for event in graph.astream_events(inputs, version=\"v2\"):\n",
    "    # 이벤트 종류와 태그 정보 추출\n",
    "    kind = event[\"event\"]\n",
    "    tags = event.get(\"tags\", [])\n",
    "\n",
    "    # 채팅 모델 스트림 이벤트 및 최종 노드 태그 필터링\n",
    "    if kind == \"on_chat_model_stream\" and \"WANT_TO_STREAM\" in tags:\n",
    "        # 이벤트 데이터 추출\n",
    "        data = event[\"data\"]\n",
    "\n",
    "        # 출력 메시지\n",
    "        if data[\"chunk\"].content:\n",
    "            print(data[\"chunk\"].content, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f100ba2",
   "metadata": {},
   "source": [
    "## 도구 호출에 대한 스트리밍 출력\n",
    "\n",
    "- `AIMessageChunk`: 실시간 토큰 단위의 출력 메시지\n",
    "- `tool_call_chunks`: 도구 호출 청크. 만약 `tool_call_chunks` 가 존재하는 경우, 도구 호출 청크를 누적하여 출력합니다. (도구 토큰은 이 속성을 보고 판단하여 출력)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17e84509",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessageChunk, HumanMessage\n",
    "\n",
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 첫 번째 메시지 처리 여부 플래그 설정\n",
    "first = True\n",
    "\n",
    "# 비동기 스트림 처리를 통한 메시지 및 메타데이터 순차 처리\n",
    "for msg, metadata in graph.stream(inputs, stream_mode=\"messages\"):\n",
    "    # 사용자 메시지가 아닌 경우의 컨텐츠 출력 처리\n",
    "    if msg.content and not isinstance(msg, HumanMessage):\n",
    "        print(msg.content, end=\"\", flush=True)\n",
    "\n",
    "    # AI 메시지 청크 처리 및 누적\n",
    "    if isinstance(msg, AIMessageChunk):\n",
    "        if first:\n",
    "            gathered = msg\n",
    "            first = False\n",
    "        else:\n",
    "            gathered = gathered + msg\n",
    "\n",
    "        # 도구 호출 청크 존재 시 누적된 도구 호출 정보 출력\n",
    "        if msg.tool_call_chunks:\n",
    "            print(gathered.tool_calls[0][\"args\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec2a8d7",
   "metadata": {},
   "source": [
    "## Subgraphs 스트리밍 출력\n",
    "\n",
    "이번에는 Subgraphs 를 통해 스트리밍 출력을 확인하는 방법을 알아보겠습니다.\n",
    "\n",
    "Subgraphs 는 그래프의 일부를 서브그래프로 정의하는 기능입니다.\n",
    "\n",
    "**흐름**\n",
    "\n",
    "- Subgraphs 에서는 기존의 최신 뉴스를 검색하는 기능을 재사용합니다.\n",
    "- Parent Graph 에서는 검색된 최신 뉴스를 바탕으로 SNS 포스트를 생성하는 기능을 추가합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "11c3b3e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, List, Dict\n",
    "from typing_extensions import TypedDict\n",
    "from langchain.tools import tool\n",
    "from langchain_teddynote.tools import GoogleNews\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ToolNode, tools_condition\n",
    "\n",
    "\n",
    "########## 1. 상태 정의 ##########\n",
    "# 상태 정의\n",
    "class State(TypedDict):\n",
    "    # 메시지 목록 주석 추가\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "\n",
    "########## 2. 도구 정의 및 바인딩 ##########\n",
    "# 도구 초기화\n",
    "# 키워드로 뉴스 검색하는 도구 생성\n",
    "@tool\n",
    "def search_keyword(query: str) -> List[Dict[str, str]]:\n",
    "    \"\"\"Look up news by keyword\"\"\"\n",
    "    news_tool = GoogleNews()\n",
    "    return \"\\n\".join(\n",
    "        [f'- {news[\"content\"]}' for news in news_tool.search_by_keyword(query, k=5)]\n",
    "    )\n",
    "\n",
    "\n",
    "# 도구 리스트 생성\n",
    "tools = [search_keyword]\n",
    "\n",
    "# LLM 초기화\n",
    "llm = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# 도구와 LLM 결합 (tags 추가)\n",
    "llm_with_tools = llm.bind_tools(tools).with_config(tags=[\"WANT_TO_STREAM\"])\n",
    "\n",
    "\n",
    "########## 3. 노드 추가 ##########\n",
    "# 챗봇 함수 정의\n",
    "def chatbot(state: State):\n",
    "    # 메시지 호출 및 반환\n",
    "    return {\"messages\": [llm_with_tools.invoke(state[\"messages\"])]}\n",
    "\n",
    "\n",
    "# SNS 포스트 생성 함수 정의\n",
    "def create_sns_post(state: State):\n",
    "    # SNS 포스트 생성을 위한 프롬프트\n",
    "    sns_prompt = \"\"\"\n",
    "    이전 대화 내용을 바탕으로 SNS 게시글 형식으로 변환해주세요.\n",
    "    다음 형식을 따라주세요:\n",
    "    - 해시태그 포함\n",
    "    - 이모지 사용\n",
    "    - 간결하고 흥미로운 문체 사용\n",
    "    - 200자 이내로 작성\n",
    "    \"\"\"\n",
    "    messages = state[\"messages\"] + [(\"human\", sns_prompt)]\n",
    "    sns_llm = ChatOpenAI(model=\"gpt-4o-mini\").with_config(tags=[\"WANT_TO_STREAM2\"])\n",
    "    return {\"messages\": [sns_llm.invoke(messages)]}\n",
    "\n",
    "\n",
    "# 서브그래프 생성\n",
    "def create_subgraph():\n",
    "    # 서브그래프용 상태 그래프 생성\n",
    "    subgraph = StateGraph(State)\n",
    "\n",
    "    # 챗봇 노드 추가\n",
    "    subgraph.add_node(\"chatbot\", chatbot)\n",
    "\n",
    "    # 도구 노드 생성 및 추가\n",
    "    tool_node = ToolNode(tools=[search_keyword])\n",
    "    subgraph.add_node(\"tools\", tool_node)\n",
    "\n",
    "    # 조건부 엣지 추가\n",
    "    subgraph.add_conditional_edges(\n",
    "        \"chatbot\",\n",
    "        tools_condition,\n",
    "    )\n",
    "\n",
    "    # tools > chatbot\n",
    "    subgraph.add_edge(\"tools\", \"chatbot\")\n",
    "\n",
    "    # START > chatbot\n",
    "    subgraph.add_edge(START, \"chatbot\")\n",
    "\n",
    "    # chatbot > END\n",
    "    subgraph.add_edge(\"chatbot\", END)\n",
    "\n",
    "    return subgraph.compile()\n",
    "\n",
    "\n",
    "# 메인 그래프 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 서브그래프 추가\n",
    "subgraph = create_subgraph()\n",
    "graph_builder.add_node(\"news_subgraph\", subgraph)\n",
    "\n",
    "# SNS 포스트 생성 노드 추가\n",
    "graph_builder.add_node(\"sns_post\", create_sns_post)\n",
    "\n",
    "# START > news_subgraph\n",
    "graph_builder.add_edge(START, \"news_subgraph\")\n",
    "\n",
    "# news_subgraph > sns_post\n",
    "graph_builder.add_edge(\"news_subgraph\", \"sns_post\")\n",
    "\n",
    "# sns_post > END\n",
    "graph_builder.add_edge(\"sns_post\", END)\n",
    "\n",
    "# 그래프 컴파일\n",
    "graph = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61fedc11",
   "metadata": {},
   "source": [
    "그래프를 시각화합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26008adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 그래프 시각화\n",
    "visualize_graph(graph, xray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a190b53",
   "metadata": {},
   "source": [
    "### Subgraphs 출력을 '생략' 하는 경우\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f2a2722",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 질문 입력\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "# 노드 업데이트 정보 순차적 처리 및 출력\n",
    "for chunk in graph.stream(inputs, stream_mode=\"updates\"):\n",
    "    # node_name: 현재 처리 중인 노드명, node_chunk: 해당 노드의 청크 데이터\n",
    "    for node_name, node_chunk in chunk.items():\n",
    "        # 현재 처리 중인 노드 구분선 출력\n",
    "        print(f\"\\n========= Update from node {node_name} =========\\n\")\n",
    "        # 해당 노드의 업데이트된 데이터 출력\n",
    "        if \"messages\" in node_chunk:\n",
    "            node_chunk[\"messages\"][-1].pretty_print()\n",
    "        else:\n",
    "            print(node_chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdeb0cbd",
   "metadata": {},
   "source": [
    "### Subgraphs 출력도 '포함' 하는 경우\n",
    "\n",
    "**참고**\n",
    "\n",
    "- `subgraphs=True` 를 통해 Subgraphs 의 출력도 포함할 수 있습니다.\n",
    "- `(namespace, chunk)` 형태로 출력됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbdffaba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 사용자의 메시지를 딕셔너리 형태로 입력 데이터 구성\n",
    "inputs = {\"messages\": [(\"human\", \"AI 관련된 최신 뉴스를 검색해줘\")]}\n",
    "\n",
    "\n",
    "# 네임스페이스 문자열을 보기 좋은 형식으로 변환하는 포맷팅 함수\n",
    "def format_namespace(namespace):\n",
    "    return namespace[-1].split(\":\")[0] if len(namespace) > 0 else \"parent graph\"\n",
    "\n",
    "\n",
    "# subgraphs=True 를 통해 서브그래프의 출력도 포함(namespace, chunk) 형태로 출력됩니다.\n",
    "for namespace, chunk in graph.stream(inputs, stream_mode=\"updates\", subgraphs=True):\n",
    "    # node_name: 현재 처리 중인 노드명, node_chunk: 해당 노드의 청크 데이터\n",
    "    for node_name, node_chunk in chunk.items():\n",
    "        print(\n",
    "            f\"\\n========= Update from node [{node_name}] in [{format_namespace(namespace)}] =========\\n\"\n",
    "        )\n",
    "\n",
    "        # 노드의 청크 데이터 출력\n",
    "        if \"messages\" in node_chunk:\n",
    "            node_chunk[\"messages\"][-1].pretty_print()\n",
    "        else:\n",
    "            print(node_chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a79a4ce4",
   "metadata": {},
   "source": [
    "#### Subgraphs 안에서 LLM 출력 토큰 단위 스트리밍"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e4fed4e",
   "metadata": {},
   "source": [
    "**참고**\n",
    "\n",
    "- `kind` 는 이벤트 종류를 나타냅니다.\n",
    "- 이벤트 종류는 [StreamEvent 타입별 정리](https://wikidocs.net/265576) 에서 확인하세요!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1bb601f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 네임스페이스 정보를 파싱하는 함수\n",
    "def parse_namespace_info(info: tuple) -> tuple[str, str]:\n",
    "    if len(info) > 1:\n",
    "        namespace, node_name = info\n",
    "        return node_name.split(\":\")[0], namespace.split(\":\")[0]\n",
    "    return info[0].split(\":\")[0], \"parent graph\"\n",
    "\n",
    "\n",
    "kind = None\n",
    "\n",
    "async for event in graph.astream_events(inputs, version=\"v2\", subgraphs=True):\n",
    "    kind = event[\"event\"]\n",
    "\n",
    "    # 이벤트 종류와 태그 정보 추출\n",
    "    if kind == \"on_chat_model_start\":\n",
    "        print(f\"\\n========= on_chat_model_start =========\\n\")\n",
    "\n",
    "    # 채팅 모델 스트림 이벤트 및 최종 노드 태그 필터링\n",
    "    elif kind == \"on_chat_model_stream\":\n",
    "        # 이벤트 데이터 추출\n",
    "        data = event[\"data\"]\n",
    "\n",
    "        # 토큰 단위의 스트리밍 출력\n",
    "        if data[\"chunk\"].content:\n",
    "            print(data[\"chunk\"].content, end=\"\", flush=True)\n",
    "\n",
    "    elif kind == \"on_tool_start\":\n",
    "        print(f\"\\n========= tool_start =========\\n\")\n",
    "        data = event[\"data\"]\n",
    "        if \"input\" in data:\n",
    "            tool_msg = data[\"input\"]\n",
    "            print(tool_msg)\n",
    "\n",
    "    elif kind == \"on_tool_end\":\n",
    "        print(f\"\\n========= tool_end =========\\n\")\n",
    "        data = event[\"data\"]\n",
    "        if \"output\" in data:\n",
    "            tool_msg = data[\"output\"]\n",
    "            print(tool_msg.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e76daf",
   "metadata": {},
   "source": [
    "#### 특정 tags 만 스트리밍 출력하는 경우\n",
    "\n",
    "- `ONLY_STREAM_TAGS` 를 통해 스트리밍 출력하고 싶은 tags 만 설정할 수 있습니다.\n",
    "- 여기서는 \"WANT_TO_STREAM\" 는 출력에서 배제하고 \"WANT_TO_STREAM2\" 만 출력하는 경우를 확인합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8305f229",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 네임스페이스 정보를 파싱하는 함수\n",
    "def parse_namespace_info(info: tuple) -> tuple[str, str]:\n",
    "    if len(info) > 1:\n",
    "        namespace, node_name = info\n",
    "        return node_name.split(\":\")[0], namespace.split(\":\")[0]\n",
    "    return info[0].split(\":\")[0], \"parent graph\"\n",
    "\n",
    "\n",
    "# 스트리밍 출력하고 싶은 tags 만 설정 (여기서는 \"WANT_TO_STREAM\" 는 출력에서 배제)\n",
    "ONLY_STREAM_TAGS = [\"WANT_TO_STREAM\"]\n",
    "\n",
    "kind = None\n",
    "tags = None\n",
    "\n",
    "async for event in graph.astream_events(inputs, version=\"v2\", subgraphs=True):\n",
    "    kind = event[\"event\"]\n",
    "    tags = event.get(\"tags\", [])\n",
    "\n",
    "    # 이벤트 종류와 태그 정보 추출\n",
    "    if kind == \"on_chat_model_start\":\n",
    "        print(f\"\\n========= tags: {tags} =========\\n\")\n",
    "\n",
    "    # 채팅 모델 스트림 이벤트 및 최종 노드 태그 필터링\n",
    "    elif kind == \"on_chat_model_stream\":\n",
    "        for tag in tags:\n",
    "            if tag in ONLY_STREAM_TAGS:\n",
    "                # 이벤트 데이터 추출\n",
    "                data = event[\"data\"]\n",
    "\n",
    "                # 출력 메시지\n",
    "                if data[\"chunk\"].content:\n",
    "                    print(data[\"chunk\"].content, end=\"\", flush=True)\n",
    "    elif kind == \"on_tool_start\":\n",
    "        print(f\"\\n========= tool_start =========\\n\")\n",
    "        data = event[\"data\"]\n",
    "        if \"input\" in data:\n",
    "            tool_msg = data[\"input\"]\n",
    "            print(tool_msg)\n",
    "\n",
    "    elif kind == \"on_tool_end\":\n",
    "        print(f\"\\n========= tool_end =========\\n\")\n",
    "        data = event[\"data\"]\n",
    "        if \"output\" in data:\n",
    "            tool_msg = data[\"output\"]\n",
    "            print(tool_msg.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
